import os
import json
from transformers import AutoTokenizer
from tqdm import tqdm

# 配置模型名称和路径
MODEL_NAME = "Qwen/Qwen-1_8B-Base"  # QWen-3-4B-Base 的 tokenizer 与此兼容
CACHE_DIR = "./model_cache"  # 缓存路径，避免重复下载

# 初始化 tokenizer（加载一次可重复使用）
def init_tokenizer():
    print("正在加载 QWen tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME,
        cache_dir=CACHE_DIR,
        trust_remote_code=True  # QWen 需要此参数
    )
    print("Tokenizer 加载完成！")
    return tokenizer

# 处理单个 JSONL 文件
def count_tokens_in_file(file_path, tokenizer, text_field="text"):
    """
    计算单个 JSONL 文件的总 token 数量
    
    参数:
    file_path: JSONL 文件路径
    tokenizer: 预加载的 tokenizer
    text_field: JSON 对象中包含文本的字段名
    
    返回:
    total_tokens: 文件的总 token 数
    details: 包含每个文本长度详情的列表
    """
    total_tokens = 0
    details = []
    error_lines = 0
    
    # 获取文件大小用于进度条
    file_size = os.path.getsize(file_path)
    
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            # 创建进度条
            with tqdm(total=file_size, unit='B', unit_scale=True, desc=f"处理 {os.path.basename(file_path)}") as pbar:
                for line in f:
                    try:
                        # 更新进度条
                        pbar.update(len(line.encode('utf-8')))
                        
                        # 解析 JSON
                        data = json.loads(line.strip())
                        
                        # 获取文本内容
                        text = data.get(text_field, "")
                        if not text:
                            continue
                            
                        # Tokenize 并计数
                        tokens = tokenizer.tokenize(text)
                        token_count = len(tokens)
                        
                        # 添加字符数供参考
                        char_count = len(text)
                        
                        # 更新计数器
                        total_tokens += token_count
                        details.append({
                            "characters": char_count,
                            "tokens": token_count,
                            "ratio": token_count / char_count if char_count else 0
                        })
                        
                    except json.JSONDecodeError:
                        error_lines += 1
                        continue
                    except Exception as e:
                        print(f"\n处理文件 {file_path} 时出错: {str(e)}")
                        error_lines += 1
                        
    except Exception as e:
        print(f"处理文件 {file_path} 时发生严重错误: {str(e)}")
        return 0, []
    
    # 显示错误行数统计
    if error_lines > 0:
        print(f"警告: {file_path} 中有 {error_lines} 行处理失败")
    
    return total_tokens, details

# 处理多个文件
def count_tokens_in_directory(directory, tokenizer, text_field="text", max_files=None):
    """
    处理目录下的所有 JSONL 文件
    
    参数:
    directory: 包含 JSONL 文件的目录
    tokenizer: 预加载的 tokenizer
    text_field: JSON 对象中包含文本的字段名
    max_files: 最多处理文件数（测试时使用）
    
    返回:
    results: 包含每个文件结果的字典
    """
    # 找到所有 JSONL 文件
    jsonl_files = [f for f in os.listdir(directory) if f.endswith('.jsonl')]
    if not jsonl_files:
        print(f"在目录 {directory} 中未找到 JSONL 文件")
        return {}
    
    # 限制处理文件数
    if max_files is not None:
        jsonl_files = jsonl_files[:max_files]
        print(f"注意: 将仅处理前 {max_files} 个文件")
    
    results = {}
    
    for filename in jsonl_files:
        file_path = os.path.join(directory, filename)
        print(f"\n开始处理文件: {filename}")
        
        total_tokens, details = count_tokens_in_file(file_path, tokenizer, text_field)
        
        # 计算字符数量（用于参考）
        char_count = sum(item['characters'] for item in details)
        
        # 计算平均长度
        avg_tokens_per_item = total_tokens / len(details) if details else 0
        
        # 保存结果
        results[filename] = {
            "total_tokens": total_tokens,
            "total_items": len(details),
            "average_tokens": avg_tokens_per_item,
            "total_characters": char_count,
            "token_to_char_ratio": total_tokens / char_count if char_count else 0,
            "details": details
        }
        
        print(f"✅ 完成处理: {filename}")
        print(f"  总 token 数: {total_tokens:,}")
        print(f"  文本条目数: {len(details):,}")
        print(f"  平均 token/文本: {avg_tokens_per_item:.2f}")
        print(f"  token/字符比率: {total_tokens / char_count if char_count else 0:.4f}")
    
    return results
#分布分析
def analyze_distribution(results):
    all_lengths = []
    for file_result in results.values():
        for detail in file_result["details"]:
            all_lengths.append(detail["tokens"])
    
    print("\nToken 长度分布分析:")
    print(f"  最大长度: {max(all_lengths)}")
    print(f"  最小长度: {min(all_lengths)}")
    print(f"  平均长度: {sum(all_lengths)/len(all_lengths):.2f}")
    print(f"  中位数: {sorted(all_lengths)[len(all_lengths)//2]}")
    
    # 长度分布直方图
    import matplotlib.pyplot as plt
    plt.hist(all_lengths, bins=50)
    plt.title("Token 长度分布")
    plt.xlabel("Token 数量")
    plt.ylabel("出现频率")
    plt.savefig("token_length_distribution.png")
    print("\n长度分布图已保存为 token_length_distribution.png")
#文本预览
def preview_texts(details, count=5):
    print("\n文本预览:")
    for i, item in enumerate(details[:count]):
        text_preview = item["text"][:70] + "..." if len(item["text"]) > 70 else item["text"]
        print(f"{i+1}. {text_preview} (Tokens: {item['tokens']})")
# 主函数
def main():
    # 初始化 tokenizer
    tokenizer = init_tokenizer()
    
    # 配置参数
    DATA_DIR = "./data"  # JSONL 文件所在目录
    TEXT_FIELD = "text"  # JSON 对象中包含文本的字段名
    
    # 处理所有文件
    results = count_tokens_in_directory(DATA_DIR, tokenizer, text_field=TEXT_FIELD)
    
    # 打印总统计
    if results:
        total_tokens = sum(r["total_tokens"] for r in results.values())
        total_items = sum(r["total_items"] for r in results.values())
        total_chars = sum(r["total_characters"] for r in results.values())
        
        print("\n" + "="*50)
        print("✨ 处理完成！汇总统计：")
        print(f"  处理的文件总数: {len(results)}")
        print(f"  所有文件的 token 总数: {total_tokens:,}")
        print(f"  所有文件的文本条目总数: {total_items:,}")
        print(f"  平均 token/条目: {total_tokens/total_items:.2f}" if total_items else "N/A")
        print(f"  总字符数: {total_chars:,}")
        print(f"  整体 token/字符比率: {total_tokens/total_chars:.4f}" if total_chars else "N/A")
        print("="*50)
        
        # 保存结果到文件
        with open("token_count_results.json", "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
            print("\n结果已保存到 token_count_results.json")

if __name__ == "__main__":
    main()